package org.example.process;

import org.apache.commons.lang3.time.DateFormatUtils;
import org.apache.flink.api.common.eventtime.TimestampAssigner;
import org.apache.flink.api.common.eventtime.TimestampAssignerSupplier;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.functions.AggregateFunction;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.api.java.tuple.Tuple3;
import org.apache.flink.streaming.api.TimerService;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.KeyedProcessFunction;
import org.apache.flink.streaming.api.functions.windowing.ProcessWindowFunction;
import org.apache.flink.streaming.api.windowing.assigners.SlidingEventTimeWindows;
import org.apache.flink.streaming.api.windowing.time.Time;
import org.apache.flink.streaming.api.windowing.windows.TimeWindow;
import org.apache.flink.util.Collector;
import org.example.flink.pojo.WaterSensor;

import java.time.Duration;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @Author: tang
 * @Description: 统计 vc 出现的次数，，获取 top n
 * @Date 2025/2/24 11:15
 */
public class TopNKeyedProcessFunction {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment();

        environment.setParallelism(1);

        SingleOutputStreamOperator<WaterSensor> streamOperator = environment.fromElements(
                new WaterSensor("1", "11", 1),
                new WaterSensor("2", "22", 2),
                new WaterSensor("3", "11", 3),
                new WaterSensor("4", "22", 4),
                new WaterSensor("4", "22", 5),
                new WaterSensor("5", "22", 6),
                new WaterSensor("5", "11", 8),
                new WaterSensor("5", "11", 10),
                new WaterSensor("5", "10", 11),
                new WaterSensor("5", "13", 13),
                new WaterSensor("5", "14", 14),
                new WaterSensor("5", "16", 16),
                new WaterSensor("5", "14", 14),
                new WaterSensor("5", "18", 18),
                new WaterSensor("5", "18", 19),
                new WaterSensor("5", "21", 22)
        ).assignTimestampsAndWatermarks(
                WatermarkStrategy.<WaterSensor>forBoundedOutOfOrderness(Duration.ofSeconds(2))
                        .withTimestampAssigner((TimestampAssignerSupplier<WaterSensor>)
                                context -> (element, recordTimestamp) -> element.getVc() * 1000L)
        );

        SingleOutputStreamOperator<Tuple3<String, Integer, Long>> aggregate = streamOperator
                .keyBy(WaterSensor::getVt)
                .window(
                        SlidingEventTimeWindows.of(Time.seconds(10), Time.seconds(5))
                )
                .aggregate(new AggregateFunction<WaterSensor, Integer, Integer>() {
                    @Override
                    public Integer createAccumulator() {
                        return 0;
                    }

                    /**
                     * 增量聚合
                     */
                    @Override
                    public Integer add(WaterSensor value, Integer accumulator) {
                        return accumulator + 1;
                    }

                    @Override
                    public Integer getResult(Integer accumulator) {
                        return accumulator;
                    }

                    @Override
                    public Integer merge(Integer a, Integer b) {
                        return null;
                    }
                }, new ProcessWindowFunction<Integer, Tuple3<String, Integer, Long>, String, TimeWindow>() {
                    // 开窗聚合后，就是普通的流，没有了窗口信息，需要自己打上窗口的标记 windowEnd
                    @Override
                    public void process(String key,
                                        ProcessWindowFunction<Integer, Tuple3<String, Integer, Long>, String, TimeWindow>.Context context,
                                        Iterable<Integer> elements,
                                        Collector<Tuple3<String, Integer, Long>> out) throws Exception {
                        // 经过 增量聚合，全窗口函数  elements 只会出现一条总的数据
                        Tuple3<String, Integer, Long> tuple3 =
                                Tuple3.of(key, elements.iterator().next(), context.window().getEnd());
                        out.collect(tuple3);
                    }
                });
        // 根据时间戳 分组
        aggregate.keyBy(tuple3 -> tuple3.f2)
                .process(new KeyedProcessFunction<Long, Tuple3<String, Integer, Long>, String>() {

                    private final Map<Long, List<Tuple2<String, Integer>>> dataListMap = new HashMap<>();

                    @Override
                    public void processElement(Tuple3<String, Integer, Long> value,
                                               KeyedProcessFunction<Long, Tuple3<String, Integer, Long>, String>.Context ctx,
                                               Collector<String> out) throws Exception {
                        TimerService timerService = ctx.timerService();
                        Tuple2<String, Integer> tuple2 = Tuple2.of(value.f0, value.f1);
                        if (!dataListMap.containsKey(value.f2)) {
                            ArrayList<Tuple2<String, Integer>> tuple2ArrayList = new ArrayList<>();
                            tuple2ArrayList.add(tuple2);
                            dataListMap.put(value.f2, tuple2ArrayList);
                        } else {
                            dataListMap.get(value.f2).add(tuple2);
                        }
//                        System.out.println("time key = " + value.f2 + "tuple2 = " + tuple2);
                        // 2. 注册一个定时器， windowEnd+1ms即可
                        // 同一个窗口范围，应该同时输出，只不过是一条一条调用processElement方法，只需要延迟1ms即可
                        timerService.registerEventTimeTimer(value.f2 + 10L);
                    }

                    @Override
                    public void onTimer(long timestamp,
                                        KeyedProcessFunction<Long, Tuple3<String, Integer, Long>, String>.OnTimerContext ctx,
                                        Collector<String> out) throws Exception {
                        super.onTimer(timestamp, ctx, out);

                        Long windowEnd = ctx.getCurrentKey();
                        List<Tuple2<String, Integer>> tuple2s = dataListMap.get(windowEnd);

//                        System.out.println("windowEnd = " + windowEnd + " tuple2s = " + tuple2s);

                        // 降序 后 - 前
                        tuple2s.sort((o1, o2) -> o2.f1 - o1.f1);

                        StringBuilder outStr = new StringBuilder();
                        outStr.append("窗口时间： ")
                                .append(windowEnd)
                                .append("\n");
                        for (int i = 0; i < Math.min(tuple2s.size(), 2); i++) {
                            Tuple2<String, Integer> tuple2 = tuple2s.get(i);
                            outStr.append("Top ").append(i + 1);
                            outStr.append(" key=").append(tuple2.f0).append(",");
                            outStr.append("count=").append(tuple2.f1).append("\n");
                        }
                        outStr.append("结束时间： ")
                                .append(windowEnd)
                                .append("\n");
                        outStr.append("================================\n");
                        out.collect(outStr.toString());

                        tuple2s.clear();
                    }

                }).print();

        environment.execute();
    }

}
